Random pick index [Reservoir Sampling]

Time: O(N); Space: O(1); medium

Given an array of integers with possible duplicates, randomly output the index of a given target number.

You can assume that the given target number must exist in the array.

Note:

  • The array size can be very large.

  • Solution that uses too much extra space will not pass the judge.

Example 1:

Input: nums=[1,2,3,3,3], target=3

Output: either index 2, 3, or 4 randomly

Explanation:

  • Each index should have equal probability of returning.

Example 2:

Input: nums=[1,2,3,3,3], target=1

Output: 0

Explanation:

  • Since in the array only nums[0] is equal to 1.

[1]:

from random import randint

class Solution1(object):

    def __init__(self, nums):
        """
        :type nums: List[int]
        """
        self.__nums = nums

    def pick(self, target):
        """
        :type target: int
        :rtype: int
        """
        reservoir = -1
        n = 0
        for i in range(len(self.__nums)):
            if self.__nums[i] != target:
                continue
            reservoir = i if randint(1, n+1) == 1 else reservoir
            n += 1
        return reservoir

[2]:

nums = [1,2,3,3,3]
s = Solution1(nums)
target = 3
assert s.pick(target) == 2 or 3 or 4
target = 1
assert s.pick(target) == 0